Pythonのreduceと内包表記／ジェネレータ式を比較してみた

こんにちは。データアナリティクス事業本部の松村です。

シーケンス（シーケンス型でなく、要素の列という意味で使っています。以下も同様です。）に対する処理を行う高階関数の代表格であるmap、filter、reduce（言語によってはfoldという名前だったりします）ですが、Pythonでは内包表記やジェネレータ式を使って同等のことができます。
mapやfilterについては内包表記／ジェネレータ式との性能差を検証した記事がたくさん見つかりますが、reduceについて言及したものは見当たりません。というわけで比べてみました。

内包表記／ジェネレータ式でreduceする

reduceと比較するにあたって、そもそも同等の処理をどうやって内包表記やジェネレータ式で書くのか、という疑問があるかもしれませんので、まずはそこから紹介します。

sumと組み合わせる

reduceの使い方として一番多いのは、シーケンスの各要素に対して何らかの計算を適用し、その結果を合算する、というものだと思います。そのような場合に限られますが、一旦計算結果をiterableで返し、最後に組み込み関数のsumで合算する、という方法が使えます。
例として、1から10までの数列について、それぞれの平方の総和を出力する処理を書いてみましょう。

sequence = range(1, 11)
sum_square = sum(x ** 2 for x in sequence)
print(sum_square)

内包表記ではなくジェネレータ式を使うのがポイントです。内包表記では平方した数列を一旦リストなどの実体に格納してしまうので、余分にメモリを消費してしまいます。

代入式（セイウチ演算子）を使う

最初に書いてしまいますが、この方法はおすすめしません。あくまで言語仕様上はこういう書き方もできなくはないというだけで、実際は素直にreduceかforループを使ったほうが良いです。

Python3.8以降であれば、新しい機能である代入式を活用することができます。代入式とは簡単に言うと、代入を行うと同時にその値を返す、という構文で、:=演算子（縦にするとセイウチの目と牙に見えるのでセイウチ演算子とも呼ばれるとか）を用います。
実際に例を挙げると、こんな感じです。意味のないコードですが...。

a = 0

def dummy():
    global a
    return (a := 1)

b = dummy()

これを実行すると、a、bともに1が代入されます。普通の代入演算子=がPythonの:=演算子相当である言語も少なくないですよね。

さて、この:=演算子と内包表記／ジェネレータ式を使ってどうreduce相当の処理を実現するかというと、こうなります。
今度は、1から10までの数列について、それぞれの平方の総乗を出力します。この方法なら総和でない集計も可能です。

import collections

sequence = range(1, 11)
product_square = sequence[0] ** 2
agg = (product_square := product_square * x ** 2 for x in sequence[1:])
collections.deque(agg, maxlen=1)
print(product_square)

まずは要素の平方値の累積を逐次product_squareに代入しつつ、その2番目からの数列をシーケンス(4,36,576,...)として生成しています。シーケンスの最後の要素と、シーケンスが最後に達したときのproduct_squareがともに求めたい値になります。
ジェネレータ式のため遅延評価されますので、シーケンスを最後まで回すためにcollections.dequeを実行しています。なぜこんなやり方をしているかというと、stackoverflow.comで、シーケンスを最後まで回すための最速の方法として紹介されていたからです。

もっとも、こんなコメントが付いていますが。

This is actually the fastest way to exhaust a long sequence, though only slighlty faster than the for loop.

+1 for being technically correct, but readers should have the usual Python caveats of, "Do you REALLY need to optimize this?", "This is less explicit, which is not Pythonic", and "The faster speed depends on implementation, which may change."

速いといってもforループよりわずかに速いだけだし、直感的でないデメリットのほうが上回る、という感じでしょうか。
もしも将来的に、シーケンスの最後の要素を取得してくれるlastなんていう関数が追加されたりすると、直感的でいい感じになりそうです。

実は似たようなことをもう少し速く処理する方法があります。さらにトリッキーなので実戦投入は絶対にやめるべきですが、あくまで参考として。

sequence = range(1, 11)
product_square = sequence[0] ** 2
[x for x in sequence[1:] if not (product_square := product_square * x ** 2)]
print(product_square)

シーケンスをひととおり回すことに関しては、collections.dequeで最後の要素を取るよりも内包表記による列挙のほうがさらに少しだけ速いのでそうしています。ただし内包表記だとメモリが余分に消費されてしまうので、if句を追加してリストへの出力を抑制しています。最終的な結果はproduct_squareにも入るので問題ありません。
ただし↓これだと期待通りになりません。

[product_square := product_square * x ** 2 for x in sequence[1:] if False]

こうするとリスト出力が抑制されるだけでなく、そもそも計算が一切行われなくなってしまいます。なので計算処理をifの条件式に移しています。
これならシーケンスの全ての要素に対して計算が行われます。途中の計算結果がFalseとして評価される値（数値の場合は0）だとそのときの要素の値がリストに出力されてしまいますが、そういうケースは稀なのでメモリ効率としては無視できる範囲です。

でも繰り返しますが、ちょっと速いからといってこんなコードを実際の製品／サービスに残すのはやめましょうね。

forループやreduceとの比較

では集計処理をいろいろな方法で行った場合のパフォーマンスを比較してみましょう。お題は1から1,000,000までの数列について、それぞれの平方の総和を求めます。
こんなコードで計測します。

from time import perf_counter
from functools import reduce

sequence = range(1, 1000001)
start_time = perf_counter()

# begin 実際の処理
# end 実際の処理

print(f"sum_square:{sum_square}")
print(f"duration:{perf_counter() - start_time}sec.")

実行環境はMacBook Pro (13-inch, 2019, Four Thunderbolt 3 ports)／第8世代クアッドコア Core i5 2.4GHz／Python3.8.2で、3回実行します。
エントリーはこちらです。

• forループ
• reduce
• ジェネレータ式とsum
• ジェネレータ式と:=演算子
• 内包表記と:=演算子

これらのコードは実際の処理部分だけ記載していきます。

forループ

sum_square = 0
for x in sequence:
    sum_square += x ** 2

duration:0.367758261sec.
duration:0.373196675sec.
duration:0.377323157sec.

reduce

sum_square = reduce(lambda acc, cur: acc + cur ** 2, sequence, 0)

duration:0.35668015200000003sec.
duration:0.34833284000000003sec.
duration:0.355396472sec.

ジェネレータ式とsum

sum_square = sum(x ** 2 for x in sequence)

duration:0.306952117sec.
duration:0.305679504sec.
duration:0.30558702400000004sec.

ジェネレータ式と:=演算子

import collections

sum_square = 0
agg = (sum_square := sum_square + x ** 2 for x in sequence)
collections.deque(agg, maxlen=1)

duration:0.36285266sec.
duration:0.371814941sec.
duration:0.371803541sec.

内包表記と:=演算子

sum_square = 0
[x for x in sequence if not (sum_square := sum_square + x ** 2)]

duration:0.337168199sec.
duration:0.34352967300000004sec.
duration:0.33432155sec.

結果比較とまとめ

速い順に概ね次のような結果になりました。

1. ジェネレータ式とsum
2. 内包表記と:=演算子
3. reduce
4. ジェネレータ式と:=演算子
5. forループ

ただし2と3、4と5の違いはそれぞれ誤差程度でしかありません。
今回のコードでは、reduceだけ要素ごとに関数呼び出しが発生するので、もっと遅いと思っていましたが意外でした。『ジェネレータ式と:=演算子』のような書き方は、たとえ今後lastのような関数が標準で使えるようになったとしても、パフォーマンスでreduceを逆転できないかもしれませんね。
そして一番遅いforループも、他の方法と比べて極端に遅いわけではありません。

というわけで、シーケンスを集計して一つの値にまとめる場合は、以下の順序で考えると良さそうです。

• 適用可能であれば、ジェネレータ式で個別の要素を計算してから最後にsumで合算する
• reduceを使う
• reduceで可読性が落ちると判断したら素直にforループ